#include "kalman_3d.h"


/**
 * 初始化三维卡尔曼滤波器（匀速度模型）
 * dt：采样时间(s)，process_noise：过程噪声系数，measure_noise：测量噪声系数
 */
void kalman3d_init(KalmanFilter3D *kf, double dt, double process_noise, double measure_noise)
{
    // 状态向量初始化（x, y, z, vx, vy, vz）
    for (int i = 0; i < 6; i++)
    {
        kf->x[i] = 0.0;
    }

    // 协方差矩阵 P 初始化（对角线为初始不确定性）
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            kf->P[i][j] = (i == j) ? 100.0 : 0.0; // 初始不确定性较大
        }
    }

    // 状态转移矩阵 F（三维匀速度模型）
    // [1 0 0 dt 0 0;
    //  0 1 0 0 dt 0;
    //  0 0 1 0 0 dt;
    //  0 0 0 1 0 0;
    //  0 0 0 0 1 0;
    //  0 0 0 0 0 1]
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            kf->F[i][j] = 0.0;
        }
    }
    kf->F[0][0] = 1;
    kf->F[0][3] = dt;
    kf->F[1][1] = 1;
    kf->F[1][4] = dt;
    kf->F[2][2] = 1;
    kf->F[2][5] = dt;
    kf->F[3][3] = 1;
    kf->F[4][4] = 1;
    kf->F[5][5] = 1;

    // 观测矩阵 H（只观测x, y, z）
    // [1 0 0 0 0 0;
    //  0 1 0 0 0 0;
    //  0 0 1 0 0 0]
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            kf->H[i][j] = (i == j) ? 1.0 : 0.0;
        }
    }

    // 过程噪声协方差 Q（与dt和过程噪声系数相关）
    double dt2 = dt * dt;
    double dt3 = dt2 * dt;
    double dt4 = dt3 * dt;
    double q = process_noise; // 过程噪声系数

    // 位置噪声分量（x, y, z对应行/列）
    kf->Q[0][0] = dt4 / 4 * q;
    kf->Q[0][3] = dt3 / 2 * q;
    kf->Q[1][1] = dt4 / 4 * q;
    kf->Q[1][4] = dt3 / 2 * q;
    kf->Q[2][2] = dt4 / 4 * q;
    kf->Q[2][5] = dt3 / 2 * q;
    // 速度噪声分量（vx, vy, vz对应行/列）
    kf->Q[3][0] = dt3 / 2 * q;
    kf->Q[3][3] = dt2 * q;
    kf->Q[4][1] = dt3 / 2 * q;
    kf->Q[4][4] = dt2 * q;
    kf->Q[5][2] = dt3 / 2 * q;
    kf->Q[5][5] = dt2 * q;

    // 测量噪声协方差 R（雷达测距/测角转换后的噪声）
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            kf->R[i][j] = (i == j) ? measure_noise : 0.0;
        }
    }
}

/**
 * 卡尔曼滤波时间更新（预测）
 */
void kalman3d_predict(KalmanFilter3D *kf)
{
    // 1. 预测状态：x_pred = F * x
    double x_pred[6];
    for (int i = 0; i < 6; i++)
    {
        x_pred[i] = 0.0;
        for (int j = 0; j < 6; j++)
        {
            x_pred[i] += kf->F[i][j] * kf->x[j];
        }
    }

    // 2. 预测协方差：P_pred = F * P * F^T + Q
    double F_T[6][6];
    mat66_transpose(kf->F, F_T); // F的转置

    double F_P[6][6];
    mat66_mult(kf->F, kf->P, F_P); // F*P

    double P_pred[6][6];
    mat66_mult(F_P, F_T, P_pred); // F*P*F^T

    // 加上过程噪声Q
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            P_pred[i][j] += kf->Q[i][j];
        }
    }

    // 更新状态和协方差
    for (int i = 0; i < 6; i++)
    {
        kf->x[i] = x_pred[i];
        for (int j = 0; j < 6; j++)
        {
            kf->P[i][j] = P_pred[i][j];
        }
    }
}

/**
 * 卡尔曼滤波量测更新（校正）
 * z_x, z_y, z_z：测量值（直角坐标）
 */
void kalman3d_update(KalmanFilter3D *kf, double z_x, double z_y, double z_z)
{
    // 1. 测量残差：y = z - H*x
    double y[3] = {
        z_x - (kf->H[0][0] * kf->x[0] + kf->H[0][1] * kf->x[1] + kf->H[0][2] * kf->x[2]),
        z_y - (kf->H[1][0] * kf->x[0] + kf->H[1][1] * kf->x[1] + kf->H[1][2] * kf->x[2]),
        z_z - (kf->H[2][0] * kf->x[0] + kf->H[2][1] * kf->x[1] + kf->H[2][2] * kf->x[2])};

    // 2. 残差协方差：S = H*P*H^T + R
    double H_T[6][3]; // H的转置（6x3）
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            H_T[i][j] = kf->H[j][i];
        }
    }

    double H_P[3][6]; // H*P（3x6）
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            H_P[i][j] = 0.0;
            for (int k = 0; k < 6; k++)
            {
                H_P[i][j] += kf->H[i][k] * kf->P[k][j];
            }
        }
    }

    double H_P_Ht[3][3]; // H*P*H^T（3x3）
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            H_P_Ht[i][j] = 0.0;
            for (int k = 0; k < 6; k++)
            {
                H_P_Ht[i][j] += H_P[i][k] * H_T[k][j];
            }
        }
    }

    // 加上测量噪声R
    double S[3][3];
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            S[i][j] = H_P_Ht[i][j] + kf->R[i][j];
        }
    }

    // 3. 卡尔曼增益：K = P*H^T*S^{-1}
    double S_inv[3][3];
    if (!mat33_inv(S, S_inv))
        return; // S不可逆，跳过更新

    double P_Ht[6][3]; // P*H^T（6x3）
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            P_Ht[i][j] = 0.0;
            for (int k = 0; k < 6; k++)
            {
                P_Ht[i][j] += kf->P[i][k] * H_T[k][j];
            }
        }
    }

    double K[6][3]; // 卡尔曼增益（6x3）
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            K[i][j] = 0.0;
            for (int k = 0; k < 3; k++)
            {
                K[i][j] += P_Ht[i][k] * S_inv[k][j];
            }
        }
    }

    // 4. 更新状态：x = x_pred + K*y
    for (int i = 0; i < 6; i++)
    {
        kf->x[i] += K[i][0] * y[0] + K[i][1] * y[1] + K[i][2] * y[2];
    }

    // 5. 更新协方差：P = (I - K*H) * P_pred
    double I[6][6] = {0}; // 6x6单位矩阵
    for (int i = 0; i < 6; i++)
        I[i][i] = 1.0;

    double K_H[6][6]; // K*H（6x6）
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            K_H[i][j] = 0.0;
            for (int k = 0; k < 3; k++)
            {
                K_H[i][j] += K[i][k] * kf->H[k][j];
            }
        }
    }

    double I_KH[6][6]; // I - K*H
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            I_KH[i][j] = I[i][j] - K_H[i][j];
        }
    }

    double P_new[6][6];
    mat66_mult(I_KH, kf->P, P_new); // (I-KH)*P_pred

    // 更新协方差矩阵
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            kf->P[i][j] = P_new[i][j];
        }
    }
}
